import java.io.*;
import java.nio.charset.Charset;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.util.Random;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.Mac;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * An ATMProtocol processes local commands sent to the ATM and writes to or reads
 * from the router as necessary. You can use whatever method you would like to
 * read from and write to the router, but this is an example to get you started.
 */

public class ATMProtocol implements Protocol {

	private PrintWriter writer;
	private BufferedReader reader;

	private static String BANK_KEY = "9F7E3D9AB71D3F127E041ACF4355F962";
	private static String BANK_IV = "683C4FCA500A1978A1A7983A0728A664";

	private int counter;
	private Integer BankCounter;

	public ATMProtocol(InputStream inputStream, OutputStream outputStream) {
		Random r = new Random();
		writer = new PrintWriter(outputStream, true);
		reader = new BufferedReader(new InputStreamReader(inputStream));
		counter = r.nextInt();
	}

	/* Continue to read input until terminated. */
	public void processLocalCommands(BufferedReader stdIn, String prompt) throws IOException {
		String userInput;

		while((userInput = stdIn.readLine()) != null) {
			processCommand(userInput, stdIn);
			counter++;
			System.out.print(prompt);
		}

		stdIn.close();
	}

	/* Interpret a command sent to the ATM and print the result to the output stream. */
	private void processCommand(String command, BufferedReader stdIn) throws IOException {
		String beginSession = "^begin-session ([A-Za-z]+)$";
		String pin = "^([0-9]{4})$";
		String balance = "^balance$";
		String withdraw = "^withdraw ([0-9]+)$";
		String endSession = "^end-session$";


		if (command.matches(balance)){
			command += "\t" + counter;
			command += "\t" + generateMAC(command, BANK_KEY);

			writer.print(encrypt(command, BANK_KEY, BANK_IV)+"\n");
			writer.flush();
			String userInput;

			while((userInput = reader.readLine()) != null){
				userInput = decrypt(userInput, BANK_KEY, BANK_IV);
				boolean tamperORreplay = true;
				
				if (checkMAC(userInput)) {
					String[] response = userInput.split("\t");
					if (BankCounter == null)
						BankCounter = new Integer(response[1]);
					
					if (BankCounter.equals(Integer.parseInt(response[1]))) {
						tamperORreplay = false;
						System.out.println(response[0]);
						BankCounter++;
					}
				} 
				
				if (tamperORreplay){
					System.out.println("Response has been tampered with.");
				}
				break;
			}

		}else if(command.matches(withdraw)){
			command += "\t" + counter;
			command += "\t" + generateMAC(command, BANK_KEY);

			writer.print(encrypt(command, BANK_KEY, BANK_IV)+"\n");
			writer.flush();
			String userInput;

			while((userInput = reader.readLine()) != null){
				userInput = decrypt(userInput, BANK_KEY, BANK_IV);
				boolean tamperORreplay = true;
				
				if (checkMAC(userInput)) {
					String[] response = userInput.split("\t");
					if (BankCounter == null)
						BankCounter = new Integer(response[1]);
					
					if (BankCounter.equals(Integer.parseInt(response[1]))) {
						tamperORreplay = false;
						System.out.println(response[0]);
						BankCounter++;
					}
				} 
				
				if (tamperORreplay){
					System.out.println("Response has been tampered with.");
				}
				break;
			}


		}else if (command.matches(beginSession)){
			String userInput;
			String userPin = "";

			System.out.print("PIN? ");
			while((userInput = stdIn.readLine()) != null)
				if (userInput.matches(pin)){
					userPin = userInput;
					break;
				}
			command += " " + userPin + "\t" + counter;
			command += "\t" + generateMAC(command, BANK_KEY);

			writer.print(encrypt(command, BANK_KEY, BANK_IV)+"\n");
			writer.flush();

			while((userInput = reader.readLine()) != null) {
				userInput = decrypt(userInput, BANK_KEY, BANK_IV);
				boolean tamperORreplay = true;
				
				if (checkMAC(userInput)) {
					String[] response = userInput.split("\t");
					if (BankCounter == null)
						BankCounter = new Integer(response[1]);
					
					if (BankCounter.equals(Integer.parseInt(response[1]))) {
						tamperORreplay = false;
						System.out.println(response[0]);
						BankCounter++;
					}
				} 
				
				if (tamperORreplay){
					System.out.println("Response has been tampered with.");
				}
				break;
			}
		}else if (command.matches(endSession)){
			String userInput;
			command += "\t" + counter;
			command += "\t" + generateMAC(command, BANK_KEY);

			writer.print(encrypt(command, BANK_KEY, BANK_IV)+"\n");
			writer.flush();

			while((userInput = reader.readLine()) != null){
				userInput = decrypt(userInput, BANK_KEY, BANK_IV);
				boolean tamperORreplay = true;
				
				if (checkMAC(userInput)) {
					String[] response = userInput.split("\t");
					if (BankCounter == null)
						BankCounter = new Integer(response[1]);
					
					if (BankCounter.equals(Integer.parseInt(response[1]))) {
						tamperORreplay = false;
						System.out.println(response[0]);
						BankCounter++;
					}
				} 
				
				if (tamperORreplay){
					System.out.println("Response has been tampered with.");
				}
				break;
			}
		}
	}

	/* Clean up all open streams. */
	public void close() throws IOException {
		writer.close();
		reader.close();
	}

	/** 
	 * Encrypts message using given key and initialization vector.
	 * 
	 * @param str plaintext message to be encrypted as a string
	 * @param keyInHex key as a hex string
	 * @param ivInHex initialization vector as a hex string
	 * @return Encrypted message as a hex string
	 */
	private String encrypt(String str, String keyInHex, String ivInHex){

		// Initialize key
		Key key = new SecretKeySpec(hexStringToByteArray(keyInHex), "AES");

		// Initialize IV
		IvParameterSpec iv = new IvParameterSpec(hexStringToByteArray(ivInHex));

		// Initialize Message
		byte[] strInBytes = str.getBytes();

		// Initialize cipher
		Cipher ciph = null;
		try {
			ciph = Cipher.getInstance("AES/CBC/PKCS5Padding");
			ciph.init(Cipher.ENCRYPT_MODE, key, iv);
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (InvalidAlgorithmParameterException e) {
			e.printStackTrace();
		}

		byte[] encryptedMessageInBytes = null;
		try {
			encryptedMessageInBytes = ciph.doFinal(strInBytes);
		} catch (IllegalBlockSizeException e) {
			e.printStackTrace();
		} catch (BadPaddingException e) {
			e.printStackTrace();
		}

		String encryptedMessage = "";
		for (int i = 0; i < encryptedMessageInBytes.length; i++) {
			encryptedMessage += String.format("%02X", encryptedMessageInBytes[i]);
		}

		return encryptedMessage;
	}

	/** 
	 * Uses HmacSHA1 to validate a message sent by the ATM.
	 * 
	 * @param str plaintext message to be encrypted as a string
	 * @param keyInHex key as a hex string
	 * @return true if message is valid or false if it has been tampered with
	 */

	private boolean checkMAC (String str) {

		if (str.equals("")) {
			return false;
		}

		String[] s = str.split("\t");
		
		if (s.length != 3) {
			return false;
		}
		
		String command = s[0]+"\t"+s[1];
		
		String checkTag = generateMAC(command, BANK_KEY);

		return checkTag.equals(s[2]);
	}

	/** 
	 * Uses HmacSHA1 to create a tag from message with given key.
	 * 
	 * @param str plaintext message to be encrypted as a string
	 * @param keyInHex key as a hex string
	 * @return MAC tag as a hex string
	 */
	private String generateMAC(String str, String keyInHex) {

		// Initialize key
		Key key = new SecretKeySpec(hexStringToByteArray(keyInHex), "HmacSHA1");

		// Initialize Message
		byte[] strInBytes = str.getBytes();

		// Initialize Mac
		Mac mac = null;
		try {
			mac = Mac.getInstance("HmacSHA1");
			mac.init(key);
		} catch (NoSuchAlgorithmException e1) {
			e1.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		}

		byte[] taggedMessageInBytes = null;
		taggedMessageInBytes = mac.doFinal(strInBytes);

		String taggedMessage = "";
		for (int i = 0; i < taggedMessageInBytes.length; i++) {
			taggedMessage += String.format("%02X", taggedMessageInBytes[i]);
		}

		return taggedMessage;
	}

	/** 
	 * Decrypts message using given key and initialization vector.
	 * 
	 * @param strInHex message to be decrypted as a hex string
	 * @param keyInHex key as a hex string
	 * @param ivInHex initialization vector as a hex string
	 * @return Decrypted message as a hex string
	 */
	private String decrypt(String strInHex, String keyInHex, String ivInHex){

		// Initialize key
		Key key = new SecretKeySpec(hexStringToByteArray(keyInHex), "AES");

		// Initialize IV
		IvParameterSpec iv = new IvParameterSpec(hexStringToByteArray(ivInHex));

		// Initialize Message
		byte[] strInBytes = hexStringToByteArray(strInHex);

		// Initialize cipher
		Cipher ciph = null;
		try {
			ciph = Cipher.getInstance("AES/CBC/PKCS5Padding");
			ciph.init(Cipher.DECRYPT_MODE, key, iv);
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		} catch (NoSuchPaddingException e) {
			e.printStackTrace();
		} catch (InvalidKeyException e) {
			e.printStackTrace();
		} catch (InvalidAlgorithmParameterException e) {
			e.printStackTrace();
		}

		byte[] decryptedMessage = null;
		try {
			decryptedMessage = ciph.doFinal(strInBytes);
		} catch (IllegalBlockSizeException e) {
			return "";
		} catch (BadPaddingException e) {
			return "";
		}

		return new String(decryptedMessage, Charset.forName("US-ASCII"));
	}

	private static byte[] hexStringToByteArray(String s) {
		int len = s.length();
		byte[] data = new byte[len / 2];
		for (int i = 0; i < len; i += 2) {
			data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
					+ Character.digit(s.charAt(i+1), 16));
		}
		return data;
	}

}
